import time
from os import path, makedirs, listdir
from fastapi import UploadFile
import paddleclas
import cv2
from config.main import options

async def save_file(file: UploadFile):
    filename = file.filename
    headers = file.headers
    type = headers.get("Content-Type")
    dir = options.get("cache_dir")
    now = time.time()
    dir = "/".join([dir, type, time.strftime("%Y%m%d", time.localtime(now))])
    dir = dir.replace("\\", "/", -1)
    if not path.exists(dir):
        makedirs(dir, 0o777)
    count = len(listdir(dir))
    arrs = filename.split(".")
    ext = arrs[len(arrs) - 1]
    filename = f"{dir}/{count+1}.{ext}"
    try:
        with open(filename, 'wb') as buffers:
            buf = await file.read(file.size)
            buffers.write(buf)
            buffers.close()
            img_mode = paddleclas.PaddleClas(model_name='text_image_orientation')
            mode_data = img_mode.predict(input_data=filename)
            model = next(mode_data)
            angle = model[0]['label_names'][0]
            cv_rotate_code = {
                '90': cv2.ROTATE_90_COUNTERCLOCKWISE,
                '180': cv2.ROTATE_180,
                '270': cv2.ROTATE_90_CLOCKWISE
            }
            if angle in cv_rotate_code:
                img = cv2.imread(filename)
                rotate_img = cv2.rotate(img, cv_rotate_code[angle])
                cv2.imwrite(filename, rotate_img)
            return filename
    except FileExistsError:
        return None